package org.wzl.service.utils;

import com.nimbusds.jose.*;
import com.nimbusds.jose.crypto.RSADecrypter;
import com.nimbusds.jose.crypto.RSAEncrypter;
import com.nimbusds.jose.crypto.RSASSASigner;
import com.nimbusds.jose.crypto.RSASSAVerifier;
import com.nimbusds.jwt.EncryptedJWT;
import com.nimbusds.jwt.JWTClaimsSet;
import com.nimbusds.jwt.SignedJWT;
import jakarta.annotation.PostConstruct;
import lombok.AllArgsConstructor;
import lombok.Data;
import lombok.ToString;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.core.io.ClassPathResource;
import org.springframework.stereotype.Component;
import org.wzl.service.exception.BusinessException;

import java.nio.file.Files;
import java.security.KeyFactory;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Base64;
import java.util.Date;
import java.util.UUID;

@Slf4j
@Component
public class JwtUtil {

    private static RSAPrivateKey privateKey;
    private static RSAPublicKey publicKey;

    @Value("${jwt.privateKeyPath}")
    private String privateKeyPath;

    @Value("${jwt.publicKeyPath}")
    private String publicKeyPath;

    @PostConstruct
    public void init() {
        try {
            privateKey = (RSAPrivateKey) loadPrivateKey(privateKeyPath);
            publicKey = (RSAPublicKey) loadPublicKey(publicKeyPath);
        } catch (Exception e) {
            log.error("Error loading keys", e);
        }
    }

    private PrivateKey loadPrivateKey(String privateKeyPath) throws Exception {
        String privateKeyPem = new String(Files.readAllBytes(new ClassPathResource(privateKeyPath).getFile().toPath()));
        privateKeyPem = privateKeyPem.replace("-----BEGIN PRIVATE KEY-----", "")
                .replace("-----END PRIVATE KEY-----", "")
                .replaceAll("\\s", "");
        byte[] keyBytes = Base64.getDecoder().decode(privateKeyPem);
        PKCS8EncodedKeySpec spec = new PKCS8EncodedKeySpec(keyBytes);
        KeyFactory kf = KeyFactory.getInstance("RSA");
        return kf.generatePrivate(spec);
    }

    private PublicKey loadPublicKey(String publicKeyPath) throws Exception {
        String publicKeyPem = new String(Files.readAllBytes(new ClassPathResource(publicKeyPath).getFile().toPath()));
        publicKeyPem = publicKeyPem.replace("-----BEGIN PUBLIC KEY-----", "")
                .replace("-----END PUBLIC KEY-----", "")
                .replaceAll("\\s", "");
        byte[] keyBytes = Base64.getDecoder().decode(publicKeyPem);
        X509EncodedKeySpec spec = new X509EncodedKeySpec(keyBytes);
        KeyFactory kf = KeyFactory.getInstance("RSA");
        return kf.generatePublic(spec);
    }

    public static String jwt(String id, String username, ExpirationTime expirationTime) {
        try {
            // Create RSA-signer with the private key
            JWSSigner signer = new RSASSASigner(privateKey);

            // Prepare JWT with claims set
            JWTClaimsSet claimsSet = new JWTClaimsSet.Builder()
                    .claim("id", id)
                    .claim("username", username)
                    .expirationTime(expirationTime.expirationTime)
                    .jwtID(UUID.randomUUID().toString())
                    .build();

            SignedJWT signedJWT = new SignedJWT(
                    new JWSHeader(JWSAlgorithm.RS256),
                    claimsSet);

            // Compute the RSA signature
            signedJWT.sign(signer);

            // Create JWE object with signed JWT as payload
            JWEObject jweObject = new JWEObject(
                    new JWEHeader.Builder(JWEAlgorithm.RSA_OAEP, EncryptionMethod.A256GCM)
                            .contentType("JWT") // required to indicate nested JWT
                            .build(),
                    new Payload(signedJWT));

            // Encrypt with the recipient's public key
            jweObject.encrypt(new RSAEncrypter(publicKey));

            // Serialize to compact JOSE form...
            return jweObject.serialize();

        } catch (Exception e) {
            log.error("Error creating encrypted JWT", e);
            return null;
        }
    }

    public static UserTokenBO parseJwt(String jwe) {
        try {
            // Parse the JWE string
            JWEObject jweObject = JWEObject.parse(jwe);

            // Decrypt with private key
            jweObject.decrypt(new RSADecrypter(privateKey));

            // Extract payload
            SignedJWT signedJWT = jweObject.getPayload().toSignedJWT();

            // Create RSA verifier
            JWSVerifier verifier = new RSASSAVerifier(publicKey);

            // Verify the signature
            if (signedJWT.verify(verifier)) {
                JWTClaimsSet claims = signedJWT.getJWTClaimsSet();
                String id = claims.getStringClaim("id");
                String username = claims.getStringClaim("username");
                return new UserTokenBO(id, username);
            } else {
                throw new RuntimeException("JWT verification failed");
            }

        } catch (Exception e) {
            throw new BusinessException("解析token失败");
        }
    }

    public static boolean isJwtExpired(String jwe) {
        try {
            // Parse the JWE string
            JWEObject jweObject = JWEObject.parse(jwe);

            // Decrypt with private key
            jweObject.decrypt(new RSADecrypter(privateKey));

            // Extract payload
            SignedJWT signedJWT = jweObject.getPayload().toSignedJWT();

            JWTClaimsSet claims = signedJWT.getJWTClaimsSet();
            Date expirationTime = claims.getExpirationTime();

            return expirationTime.before(new Date());

        } catch (Exception e) {
            log.error("Error checking if JWT is expired", e);
            return true;
        }
    }

    @AllArgsConstructor
    public enum ExpirationTime {
        NO_README(new Date(System.currentTimeMillis() + 1000 * 60 * 60)),
        README(new Date(System.currentTimeMillis() + 1000 * 60 * 60 * 24 * 7));

        private final Date expirationTime;
    }

    @Data
    @ToString
    @AllArgsConstructor
    public static class UserTokenBO {
        private String id;
        private String username;
    }
}
